#include "dyn_ode_dopri54.h"

DOPRI54::DOPRI54(MultiFunctor &func,
                 colvec _initCond,
                 double _start,
                 double _end,
                 double _initStepsize,
                 double _minStepsize,
                 double _maxStepsize,
                 double _absTol,
                 double _relTol,
                 int _maxIterations)
    : ODE(func,_initCond,_start,_end,_initStepsize,_minStepsize,_maxStepsize,_absTol,_relTol,_maxIterations)
{
}

colvec DOPRI54::next(colvec previousStep)
{
    while(1)
    {
        // We compute the k_i vectors.
        colvec k1 = initStepsize*integrand(previousStep);
        colvec k2 = initStepsize*integrand(previousStep+a21*k1);
        colvec k3 = initStepsize*integrand(previousStep+a31*k1+a32*k2);
        colvec k4 = initStepsize*integrand(previousStep+a41*k1+a42*k2+a43*k3);
        colvec k5 = initStepsize*integrand(previousStep+a51*k1+a52*k2+a53*k3+a54*k4);
        colvec k6 = initStepsize*integrand(previousStep+a61*k1+a62*k2+a63*k3+a64*k4+a65*k5);

        // We compute the fifth- and fourth-order increments.
        colvec fifthOrder = b1*k1+b2*k2+b3*k3+b4*k4+b5*k5+b6*k6;
        colvec fourthOrder = c1*k1+c2*k2+c3*k3+c4*k4+c5*k5+c6*k6;

        // We compute the difference of the methods, which gives
        // an error estimate delta. We take the user-supplied
        // variables absTol and relTol to form a scale variable.
        // We define an error estimate (W.H. Press et al, 2007. Numerical Recipes,
        // third edn, p.913).
        colvec delta = fifthOrder-fourthOrder;
        colvec scale = relTol*(previousStep+fifthOrder)+absTol;

        colvec normalizedSquaredError = pow(delta/scale,2.0);
        double error = sqrt((1./((double) normalizedSquaredError.n_rows))*accu(normalizedSquaredError));

        // We compute the stepsize for the next
        // iteration if the steps succeeds and
        // the same step if it does not.
        double scaling = pow(error,-alpha)*pow(errorOld,beta);

        // We control the change in the stepsize
        if (scaling<minScale) scaling=minScale;
        if (scaling>maxScale) scaling=maxScale;

        oldStep = initStepsize;
        initStepsize *= scaling;

        bool trigger = false;
        if (initStepsize<minStepsize){initStepsize=minStepsize;trigger=true;}
        if (initStepsize>maxStepsize){initStepsize=maxStepsize;trigger=true;}

        if (error <= 1.0 || trigger)
        {
            errorOld = error;
            previousStep += fifthOrder;
            return previousStep;
        }
    }
}

int DOPRI54::integrate()
{
    cout << "Preparing DOPRI54 integration...";
    // We create a matrix having a single row,
    // but with enough columns to store all the data.
    int dimension = initCond.n_rows;
    mat tempResults = mat(1,dimension+1);

    // The first row contains the initial data.
    tempResults(0,0) = start;
    tempResults.submat(span(0,0),span(1,dimension)) = initCond.t();

    // We now compute the complete integral.
    // Since we do not know in advance on many
    // rows the results matrix will have, we use a
    // for loop to exit if we fail to do the integral.
    cout << "Integrating." << endl;
    for (int i=1;i<maxIterations;i++)
    {
        // We add a row to the results matrix.
        tempResults.resize(i+1,dimension+1);

        // We compute the values of the new row.
        tempResults.submat(span(i,i),span(1,dimension)) = next(tempResults.submat(span(i-1,i-1),span(1,dimension)).t()).t();
        tempResults(i,0) = tempResults(i-1,0)+getOldStep();

        if (tempResults(i,0) >= end)
        {
            i=maxIterations;
        }
    }

    results = tempResults;

    return 0;
}
